package com.example.nettyrpcdemo.core;

import cn.hutool.core.lang.Pair;
import cn.hutool.extra.spring.SpringUtil;
import com.example.nettyrpcdemo.config.GameConfig;
import com.google.common.collect.Maps;
import com.google.protobuf.GeneratedMessage;
import io.netty.channel.Channel;
import java.lang.reflect.Method;
import java.util.Collection;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.TimeUnit;
import java.util.function.Consumer;
import lombok.Getter;
import lombok.extern.slf4j.Slf4j;
import msg.MsgCross;
import msg.MsgIds;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.lang.Nullable;
import org.springframework.stereotype.Controller;
import org.springframework.stereotype.Service;
import tools.MessageManager;

/**
 * 发送消息
 */
@Slf4j
@Service
public class RpcService {

    private static final int TIMEOUT_SEC = 15;

    //<serverId, rpc客户端>
    private final Map<Integer, NettyRpcClient> rpcClientMap = Maps.newConcurrentMap();

    // <回调id, 主要是等待回应的CompletableFuture>
    @Getter
    private final Map<Long, CallContext<?>> callbackId2CallContextMap = Maps.newConcurrentMap();

    private final Map<Integer, RpcCmdExecutor> msgId2CmdExecutorMap = new HashMap<>();

    // TODO 测试代码, 正式应该从zk中获取
    private final Map<Integer, Pair<String, Integer>> serverId2PairMap = Map.of(1, Pair.of("127.0.0.1", 6001), 2, Pair.of("127.0.0.1", 6002));

    @Autowired
    private GameConfig gameConfig;

    public void init() {
        Collection<Object> values = SpringUtil.getApplicationContext().getBeansWithAnnotation(Controller.class).values();
        for (Object value : values) {
            Method[] declaredMethods = value.getClass().getDeclaredMethods();
            for (Method method : declaredMethods) {
                RpcMsgReceiver rpcMsgReceiverAnnotation = method.getAnnotation(RpcMsgReceiver.class);
                if (rpcMsgReceiverAnnotation != null) {
                    Class<?> clazz = rpcMsgReceiverAnnotation.value();
                    int msgId = MessageManager.getMsgIdByMsgClazz(clazz);

                    msgId2CmdExecutorMap.put(msgId, new RpcCmdExecutor(method, value));
                }
            }
        }

        log.info("rpcService init provider method count={}", msgId2CmdExecutorMap.size());
    }

    public void callProviderMethod(Channel channel, ProtobufMessage message) {
        GameRpcMsgContext context = new GameRpcMsgContext(channel, message.getCallbackId(), message.getMessage());

        RpcCmdExecutor rpcCmdExecutor = msgId2CmdExecutorMap.get(message.getMsgId());
        if (rpcCmdExecutor == null) {
            context.sendAnswerMsg(MsgCross.ProviderException.newBuilder().build());
            log.error("服务提供者serverId={} 未监听msgId={} ", gameConfig.getServerId(), message.getMsgId());
            return;
        }

        try {
            rpcCmdExecutor.invoke(context);
        } catch (Exception e) {
            log.error("provider exec method failed!!! error={} msgId={}", e, message.getMsgId());
            // 返回给服务消费者
            context.sendAnswerMsg(MsgCross.ProviderException.newBuilder().build());
        }
    }

    public void send(int toServerId, GeneratedMessage msg) {
        NettyRpcClient client = getNettyRpcClient(toServerId);
        if (client == null) {
            log.error("send error!!! can't connect to serverId={}, msg={}", toServerId, msg);
            return;
        }
        client.send(msg, 0);
    }

    @Nullable
    public <T extends GeneratedMessage> T syncCall(int toServerId, GeneratedMessage msg, Class<T> t) throws Exception {
        NettyRpcClient client = getNettyRpcClient(toServerId);


        CompletableFuture<T> completableFuture = new CompletableFuture<>();

        CallContext<T> callContext = new CallContext<>();
        callContext.setFuture(completableFuture);

        long nextSeqId = MsgIdGenerator.nextSeqId();
        callbackId2CallContextMap.put(nextSeqId, callContext);
        if (client == null) {
            log.error("syncCall error!!!can't connect to serverId={} msg={}", toServerId, msg);
            setFailed(nextSeqId);
            return null;
        }

        client.send(msg, nextSeqId);

        T ret = completableFuture.completeOnTimeout(null, TIMEOUT_SEC, TimeUnit.SECONDS).get();
        callbackId2CallContextMap.remove(nextSeqId);

        if (ret != null) {
            int retMsgId = MessageManager.getMsgIdByMsgClazz(ret.getClass());
            if (retMsgId == MsgIds.ProviderException) {
                ret = null;
            }
        }
        return ret;
    }

    public <T extends GeneratedMessage> void asyncCall(int toServerId, GeneratedMessage msg, Class<T> t, Consumer<T> consumer) {
        NettyRpcClient client = getNettyRpcClient(toServerId);

        CompletableFuture<T> completableFuture = new CompletableFuture<>();
        CallContext<T> callContext = new CallContext<>();
        callContext.setFuture(completableFuture);

        long nextSeqId = MsgIdGenerator.nextSeqId();
        completableFuture //
                .completeOnTimeout(null, TIMEOUT_SEC, TimeUnit.SECONDS) // RPC超时
                .whenComplete((@Nullable T ret, Throwable e) -> { // 得到结果
                    callbackId2CallContextMap.remove(nextSeqId);
                    if (ret != null) {
                        int retMsgId = MessageManager.getMsgIdByMsgClazz(ret.getClass());
                        if (retMsgId == MsgIds.ProviderException) {
                            ret = null;
                        }
                    }

                    consumer.accept(ret);
                });
        callbackId2CallContextMap.put(nextSeqId, callContext);

        if (client == null) {
            log.error("asyncCall error!!! can't connect to serverId={} msg={}", toServerId, msg);
            setFailed(nextSeqId);
            return;
        }
        client.send(msg, nextSeqId);
    }

    @Nullable
    private NettyRpcClient getNettyRpcClient(int toServerId) {
        return rpcClientMap.computeIfAbsent(toServerId, k -> {
            NettyRpcClient client = new NettyRpcClient();

            Pair<String, Integer> pair = serverId2PairMap.get(toServerId);

            String ip = pair.getKey();
            Integer port = pair.getValue();

            // 从zk读取ip和端口
            try {
                client.connect(ip, port);
            } catch (Exception e) {
                return null;
            }

            return client;
        });
    }

    private void setFailed(long callbackId) {
        CallContext<?> callContext = callbackId2CallContextMap.remove(callbackId);
        if (callContext != null) {
            callContext.getFuture().complete(null);
        }
    }
}
